import os
import subprocess
import re
from PIL import Image


def get_video_duration(video_path):
    """
    获取视频的时长
    :param video_path: 视频文件路径
    :return: 视频时长（秒）
    """
    try:
        cmd = [
            'ffmpeg',
            '-i', video_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        output = result.stderr
        # 使用正则表达式匹配时长信息
        match = re.search(r"Duration: (\d{2}):(\d{2}):(\d{2})\.\d+", output)
        if match:
            hours = int(match.group(1))
            minutes = int(match.group(2))
            seconds = int(match.group(3))
            total_seconds = hours * 3600 + minutes * 60 + seconds
            return total_seconds
        else:
            print(f"未找到视频时长信息。错误信息: {output}")
            return 0
    except Exception as e:
        print(f"获取视频时长时出错: {e}")
        return 0


def generate_preview_frames(video_path, num_frames=60):
    """
    生成指定数量的预览图
    :param video_path: 视频文件路径
    :param num_frames: 预览图数量
    :return: 预览图文件名列表
    """
    duration = get_video_duration(video_path)
    if duration == 0:
        return []
    interval = duration / num_frames
    frame_paths = []
    for i in range(num_frames):
        time = i * interval
        frame_path = f"frame_{i:03d}.jpg"
        try:
            cmd = [
                'ffmpeg',
                '-ss', str(time),
                '-i', video_path,
                '-vframes', '1',
                '-q:v', '2',
                frame_path
            ]
            subprocess.run(cmd, check=True)
            frame_paths.append(frame_path)
        except Exception as e:
            print(f"生成预览图 {frame_path} 时出错: {e}")
    return frame_paths


def merge_frames(frame_paths, output_path):
    """
    合并预览图为一张矩阵图片
    :param frame_paths: 预览图文件名列表
    :param output_path: 合并后图片的保存路径
    """
    if not frame_paths:
        return
    frames = [Image.open(frame) for frame in frame_paths]
    rows = 6
    cols = 10
    frame_width = frames[0].width
    frame_height = frames[0].height
    total_width = cols * frame_width
    total_height = rows * frame_height
    merged_image = Image.new('RGB', (total_width, total_height))
    for i in range(rows):
        for j in range(cols):
            index = i * cols + j
            if index < len(frames):
                frame = frames[index]
                x_offset = j * frame_width
                y_offset = i * frame_height
                merged_image.paste(frame, (x_offset, y_offset))
    merged_image.save(output_path)
    for frame in frame_paths:
        os.remove(frame)


def process_video_directory(video_dir, output_dir):
    """
    处理指定目录下的所有视频
    :param video_dir: 视频目录
    :param output_dir: 输出目录
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for root, _, files in os.walk(video_dir):
        for file in files:
            if file.lower().endswith(('.mp4', '.avi', '.mov')):
                video_path = os.path.join(root, file)
                frame_paths = generate_preview_frames(video_path)
                output_filename = os.path.splitext(file)[0] + '_preview.jpg'
                output_path = os.path.join(output_dir, output_filename)
                merge_frames(frame_paths, output_path)


if __name__ == "__main__":
    video_directory = '.'  # 当前目录，可根据需要修改
    output_directory = 'jpeg'
    process_video_directory(video_directory, output_directory)
